import tensorflow as tf
from dataPretreat import data_pretreat


# 创建并训练模型
def get_compiled_model():
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(10, activation='relu'),
        tf.keras.layers.Dense(10, activation='relu'),
        tf.keras.layers.Dense(1, activation='sigmoid')
    ])
    model.compile(optimizer='adam',
                  loss='binary_crossentropy',
                  metrics=['accuracy'])
    return model


# 标准化数据
def process_continuous_data(mean, data):
    data = tf.cast(data, tf.float32) * 1 / (2 * mean)
    return tf.reshape(data, [-1, 1])


train_data = data_pretreat(path='../resource/KDDTrain+.csv')
test_data = data_pretreat(path='../resource/KDDTest+.csv')
print(train_data.head())
label = train_data.pop('label')
dataset = tf.data.Dataset.from_tensor_slices((train_data.values, label.values))

for feat, targ in dataset.take(5):
    print('Features: {}, Target: {}'.format(feat, targ))

train_dataset = dataset.shuffle(len(train_data)).batch(1)
model = get_compiled_model()
model.fit(dataset, epochs=1)

test_loss, test_accuracy = model.evaluate(test_data)
print('\n\nTest Loss {}, Test Accuracy {}'.format(test_loss, test_accuracy))